package com.kongxiang.cubebit.llm.service;

import com.kongxiang.cubebit.base.common.exceptions.BaseException;
import com.kongxiang.cubebit.llm.config.OpenAiRerankProperties;
import com.kongxiang.cubebit.llm.dto.ReRankRequest;
import com.kongxiang.cubebit.llm.dto.ReRankResponse;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.autoconfigure.openai.OpenAiChatProperties;
import org.springframework.ai.autoconfigure.openai.OpenAiEmbeddingProperties;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.openai.OpenAiChatModel;
import org.springframework.http.HttpStatusCode;
import org.springframework.http.MediaType;
import org.springframework.stereotype.Service;
import org.springframework.util.StringUtils;
import org.springframework.web.client.RestClient;
import org.springframework.web.reactive.function.client.WebClient;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.util.Map;

@Service
@RequiredArgsConstructor
@Slf4j
public class OpenAIService {

    private final OpenAiChatProperties openAiChatProperties;
    private final OpenAiEmbeddingProperties openAiEmbeddingProperties;
    private final OpenAiRerankProperties openAiRerankProperties;

    private final WebClient webClient;
    private final RestClient restClient;

    private final EmbeddingModel embeddingModel;
    private final OpenAiChatModel openAiChatModel;

    public ReRankResponse rerank(ReRankRequest request) {
        /*HttpResponse<String> response = Unirest.post("https://api.siliconflow.cn/v1/rerank")
  .header("Authorization", "Bearer <token>")
  .header("Content-Type", "application/json")
  .body("{\n  \"model\": \"BAAI/bge-reranker-v2-m3\",\n  \"query\": \"Apple\",\n  \"documents\": [\n    \"apple\",\n    \"banana\",\n    \"fruit\",\n    \"vegetable\"\n  ],\n  \"top_n\": 4,\n  \"return_documents\": false,\n  \"max_chunks_per_doc\": 1024,\n  \"overlap_tokens\": 80\n}")
  .asString();*/
        return restClient.post()
                .uri(openAiRerankProperties.getBaseUrl() + "/v1/rerank")
                .headers(httpHeaders -> {
                    httpHeaders.setBearerAuth(openAiRerankProperties.getApiKey());
                    httpHeaders.setContentType(MediaType.APPLICATION_JSON);
                })
                .body(Map.of(
                        "model", request.getModel() == null ? openAiRerankProperties.getModel() : request.getModel(),
                        "query", request.getQuery(),
                        "documents", request.getDocuments() ,
                        "top_n", request.getTop_n() == 0 ? 5 : request.getTop_n(),
                        "return_documents", request.isReturn_documents()
                )).retrieve()
                .onStatus(response -> response.getStatusCode().isError()).body(ReRankResponse.class);
    }

    public float[] embedding(String text) {
        // 请求json
        return embeddingModel.embed(Document.builder().text(text).build());
    }

    public Flux<String> openAiStream(String prompt) {
        // 请求json
        String json = "{\"model\":\"" + openAiChatProperties.getOptions().getModel() + "\",\"messages\":[{\"role\":\"user\",\"content\":\"" + prompt + "\"}],\"stream\":true}";
        // WebClient 调用stream大模型
        return openAiChatModel.stream(json);

    }

    public Flux<String> post(String url, String body, String baseAuthKey) {
        WebClient.RequestHeadersSpec<?> requestHeadersSpec = webClient.post()
                .uri(url)
                .contentType(MediaType.APPLICATION_JSON)
                .bodyValue(body);
        if (StringUtils.hasText(baseAuthKey)) {
            requestHeadersSpec.headers((headers) -> {
                        headers.setBearerAuth(baseAuthKey);
                    }
            );
            log.debug("baseAuthKey: {}", baseAuthKey);
        }
        return requestHeadersSpec
                .retrieve()
                .onStatus(HttpStatusCode::isError, (clientResponse -> Mono.error(new BaseException(clientResponse.statusCode().value() + "", clientResponse.toString()))))
                .bodyToFlux(String.class);
    }
}